"""
train.py 文件运行在result/MyEnv/MyEnv/mappo/check/目录下
train.py 的配置文件为config

eval.py 文件运行在result/MyEnv_model/MyEnv/mappo/check目录下
eval.py 的配置文件为config—_eval

env_core.py 的配置文件为config_env

新增ssh修改功能
"""

# !/usr/bin/env python
import os
import sys
from pathlib import Path

import numpy as np
import setproctitle
import torch

current_dir = os.path.dirname(__file__)
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)

from config_env import get_env_config
from config_train import get_config
from envs.env_wrappers import DummyVecEnv
from envs.ruler import ruler_main, show_move_list

parser_ev = get_env_config()

"""Train script for MPEs."""


def make_train_env(all_args, env_size=[]):
    def get_env_fn(rank):
        def init_env(env_size):
            # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.

            from envs.env_continuous import ContinuousActionEnv

            env = ContinuousActionEnv(env_size)

            # from envs.env_discrete import DiscreteActionEnv

            # env = DiscreteActionEnv()

            env.seed(all_args.seed + rank * 1000)
            return env

        return init_env

    return DummyVecEnv(
        [get_env_fn(i) for i in range(all_args.n_rollout_threads)], env_size
    )


def make_eval_env(all_args, env_size=[]):
    def get_env_fn(rank):
        def init_env(env_size):
            # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.
            from envs.env_continuous import ContinuousActionEnv

            env = ContinuousActionEnv(env_size)
            # from envs.env_discrete import DiscreteActionEnv
            # env = DiscreteActionEnv()
            env.seed(all_args.seed + rank * 1000)
            return env

        return init_env

    return DummyVecEnv(
        [get_env_fn(i) for i in range(all_args.n_eval_rollout_threads)], env_size
    )


def parse_args(args, parser):
    parser.add_argument(
        "--scenario_name", type=str, default="MyEnv", help="Which scenario to run on"
    )
    parser.add_argument("--num_landmarks", type=int, default=3)
    parser.add_argument(
        "--num_agents",
        type=int,
        default=parser_ev["agent_num"],
        help="number of players",
    )

    all_args = parser.parse_known_args(args)[0]

    return all_args


def Get_Run(model_dir):
    parser = get_config()
    all_args = parse_args(sys.argv[1:], parser)

    if all_args.algorithm_name == "rmappo":
        assert (
            all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy
        ), "check recurrent policy!"
    elif all_args.algorithm_name == "mappo":
        assert (
            all_args.use_recurrent_policy == False
            and all_args.use_naive_recurrent_policy == False
        ), "check recurrent policy!"
    else:
        raise NotImplementedError

    assert (
        (
            all_args.share_policy == True
            and all_args.scenario_name == "simple_speaker_listener"
        )
        == False
    ), "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    if not run_dir.exists():
        curr_run = "run1"
    else:
        exst_run_nums = [
            int(str(folder.name).split("run")[1])
            for folder in run_dir.iterdir()
            if str(folder.name).startswith("run")
        ]
        if len(exst_run_nums) == 0:
            curr_run = "run1"
        else:
            curr_run = "run%i" % (max(exst_run_nums) + 1)
    run_dir = run_dir / curr_run
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_train_env(all_args)
    eval_envs = make_eval_env(all_args) if all_args.use_eval else None
    num_agents = all_args.num_agents
    # model_dir = "D:/Reinforce/light_mappo-main/light_mappo-main/results/MyEnv/MyEnv/mappo/check/run125/models"
    all_args.model_dir = model_dir
    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from runner.shared.env_runner import EnvRunner as Runner
    else:
        from runner.separated.env_runner import EnvRunner as Runner

    runner = Runner(config)
    return runner


def main(args, env_size=[]):
    parser = get_config()
    all_args = parse_args(args, parser)

    if all_args.algorithm_name == "rmappo":
        assert (
            all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy
        ), "check recurrent policy!"
    elif all_args.algorithm_name == "mappo":
        assert (
            all_args.use_recurrent_policy == False
            and all_args.use_naive_recurrent_policy == False
        ), "check recurrent policy!"
    else:
        raise NotImplementedError

    assert (
        (
            all_args.share_policy == True
            and all_args.scenario_name == "simple_speaker_listener"
        )
        == False
    ), "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    if not run_dir.exists():
        curr_run = "run1"
        model_dir = ""
    else:
        exst_run_nums = [
            int(str(folder.name).split("run")[1])
            for folder in run_dir.iterdir()
            if str(folder.name).startswith("run")
        ]
        if len(exst_run_nums) == 0:
            curr_run = "run1"
            model_dir = ""
        else:
            curr_run = "run%i" % (max(exst_run_nums) + 1)
            model_dir = ""
            if parser_ev["is_load_model"]:
                model_dir = "run%i/models/" % (max(exst_run_nums))
                model_dir = run_dir / model_dir
    run_dir = run_dir / curr_run
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_train_env(all_args, env_size)
    eval_envs = make_eval_env(all_args, env_size) if all_args.use_eval else None
    num_agents = all_args.num_agents
    # model_dir = "D:/Reinforce/light_mappo-main/light_mappo-main/results/MyEnv/MyEnv/mappo/check/run7/models"
    print(model_dir)
    if model_dir != "":
        all_args.model_dir = model_dir
    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from runner.shared.env_runner import EnvRunner as Runner
    else:
        from runner.separated.env_runner import EnvRunner as Runner

    runner = Runner(config)
    print("runner created")
    runner.run()
    # post process
    envs.close()
    if all_args.use_eval and eval_envs is not envs:
        eval_envs.close()

    runner.writter.export_scalars_to_json(str(runner.log_dir + "/summary.json"))
    runner.writter.close()


if __name__ == "__main__":
    if parser_ev["is_Draw_now"] and parser_ev["env_type"] == "Draw":
        ruler_main()
        show_move_list()
    elif parser_ev["env_type"] == "Rand":
        for i in range(10):
            main(sys.argv[1:], ["Circle", [-i, i + 1], 10 + i])
        for i in range(10):
            main(sys.argv[1:], ["Rectangle", [-i - 10, i - 10], 10 + i, 20 - i])
        for i in range(10):
            main(sys.argv[1:], ["Circle", [-i + 1, i + 1], 5 + i])
        for i in range(10):
            main(sys.argv[1:], ["Rectangle", [-20 + i, 20 - i * 2], 5 + i * 2, 20 - i])
    else:
        for i in range(parser_ev["train_num"]):
            main(sys.argv[1:])
